import torch
import torchvision as tv
from torch import nn
import torch.utils.data.dataset as dataset
import cv2 as cv
import numpy as np
import pandas as pd
import os

class garbageDataSet(Dataset):
    def __init__(self, dataPath, listPath):
        super().__init__()
        self.root = dataPath
        list = open(listPath)
        self.imgNames = [i.split(' ')[0] for i in list.readlines()]
        self.labels = [int(i.split(' ')[1]) for i in list.readlines()]

    def __len__(self):
        folders = os.listdir(self.root)
        length = 0
        for folder in folders:
            length += len(os.listdir(self.root + "/" + folder))
        return length

    def __getitem__(self, idx):
        print(">>get it", idx)
        img = cv.imread(self.root + "/" + self.imgNames[i])
        img = cv.resize(img, (227, 227))
        return img, label
'''train_data = garbageDataSet()
        test_data = garbageDataSet()
        train_loader = Dataloader(train_data, batch_size=128, shuffle=True, num_workers=4)
        test_loader = DataLoader(test_data, batch_size=128, shuffle=True, num_workers=4)'''

class reader(Dataset):
    def __init__(self, root):
        super().__init__()
        self.root = root
        self.labels = [int(i) for i in os.listdir(root)]
        ways = [root + '/' + i for i in os.listdir(root)]
        self.names = []
        for way in ways:
            self.names += [way + '/' + i for i in os.listdir(way)]

    def __getitem__(self, idx):
        img = cv.imread(self.names[idx])
        img = cv.resize(img, (224, 224))
        label = int(self.names[idx].split('/')[2])
        Img = torch.tensor(img, dtype=torch.float).reshape((3, 224, 224))
        Label = torch.tensor(label, dtype=torch.long)
        return Img, Label

    def __len__(self):
        length = 0
        for i in os.listdir(self.root):
            length += len(os.listdir(self.root + "/" + i))
        return length


